# coding=utf-8

import json
import os
import h5py
import numpy as np


def collect_params(file_path: str):
    print("Processing {}".format(file_path))
    files = os.listdir(file_path)
    id_array = np.array([], dtype=np.int64)
    heliostat_area_array = np.array([], dtype=np.int64)
    heliostat_l_ratio_array = np.array([], dtype=np.int64)
    heliostat_w_ratio_array = np.array([], dtype=np.int64)
    heliostat_glass_thickness_array = np.array([], dtype=np.int64)
    flux_map_shape_array = np.array([], dtype=np.int64)
    distance_heliostat_receiver_array = np.array([], dtype=np.int64)
    sun_elevation_array = np.array([], dtype=np.int64)
    sun_azimuth_array = np.array([], dtype=np.int64)
    heliostat_glass_sigma_array = np.array([], dtype=np.float64)
    sun_csr_array = np.array([], dtype=np.float64)
    heliostat_glass_refractive_index_array = np.array([], dtype=np.float64)
    peak_flux_param_array = np.array([], dtype=np.float64)
    sigma_x_array = np.array([], dtype=np.float64)
    sigma_y_array = np.array([], dtype=np.float64)
    x13_array = np.array([], dtype=np.float64)
    x10_array = np.array([], dtype=np.float64)
    y11_array = np.array([], dtype=np.float64)
    y10_array = np.array([], dtype=np.float64)
    gradient_inv_array = np.array([], dtype=np.float64)

    for current_file in files:
        current_path = os.path.join(file_path, current_file)
        if os.path.isdir(current_path):
            (
                id_sub_array,
                heliostat_area_sub_array,
                heliostat_l_ratio_sub_array,
                heliostat_w_ratio_sub_array,
                heliostat_glass_thickness_sub_array,
                flux_map_shape_sub_array,
                distance_heliostat_receiver_sub_array,
                sun_elevation_sub_array,
                sun_azimuth_sub_array,
                heliostat_glass_sigma_sub_array,
                sun_csr_sub_array,
                heliostat_glass_refractive_index_sub_array,
                peak_flux_param_sub_array,
                sigma_x_sub_array,
                sigma_y_sub_array,
                x13_sub_array,
                x10_sub_array,
                y11_sub_array,
                y10_sub_array,
                gradient_inv_sub_array,
            ) = collect_params(current_path)

            id_array = np.concatenate((id_array, id_sub_array))
            heliostat_area_array = np.concatenate(
                (heliostat_area_array, heliostat_area_sub_array)
            )
            heliostat_l_ratio_array = np.concatenate(
                (heliostat_l_ratio_array, heliostat_l_ratio_sub_array)
            )
            heliostat_w_ratio_array = np.concatenate(
                (heliostat_w_ratio_array, heliostat_w_ratio_sub_array)
            )
            heliostat_glass_thickness_array = np.concatenate(
                (heliostat_glass_thickness_array, heliostat_glass_thickness_sub_array)
            )
            flux_map_shape_array = np.concatenate(
                (flux_map_shape_array, flux_map_shape_sub_array)
            )
            distance_heliostat_receiver_array = np.concatenate(
                (
                    distance_heliostat_receiver_array,
                    distance_heliostat_receiver_sub_array,
                )
            )
            sun_elevation_array = np.concatenate(
                (sun_elevation_array, sun_elevation_sub_array)
            )
            sun_azimuth_array = np.concatenate(
                (sun_azimuth_array, sun_azimuth_sub_array)
            )
            heliostat_glass_sigma_array = np.concatenate(
                (heliostat_glass_sigma_array, heliostat_glass_sigma_sub_array)
            )
            sun_csr_array = np.concatenate((sun_csr_array, sun_csr_sub_array))
            heliostat_glass_refractive_index_array = np.concatenate(
                (
                    heliostat_glass_refractive_index_array,
                    heliostat_glass_refractive_index_sub_array,
                )
            )
            peak_flux_param_array = np.concatenate(
                (peak_flux_param_array, peak_flux_param_sub_array)
            )
            sigma_x_array = np.concatenate((sigma_x_array, sigma_x_sub_array))
            sigma_y_array = np.concatenate((sigma_y_array, sigma_y_sub_array))
            x13_array = np.concatenate((x13_array, x13_sub_array))
            x10_array = np.concatenate((x10_array, x10_sub_array))
            y11_array = np.concatenate((y11_array, y11_sub_array))
            y10_array = np.concatenate((y10_array, y10_sub_array))
            gradient_inv_array = np.concatenate(
                (gradient_inv_array, gradient_inv_sub_array)
            )

        else:
            _, file_extension = os.path.splitext(current_path)
            if file_extension == ".json":
                current_parent_path, _ = os.path.split(current_path)
                current_parent_path, current_heliostat_glass_thickness_str = (
                    os.path.split(current_parent_path)
                )
                current_parent_path, current_heliostat_l_w_ratio_str = os.path.split(
                    current_parent_path
                )
                current_parent_path, current_heliostat_area_str = os.path.split(
                    current_parent_path
                )
                # 单位：毫米
                current_heliostat_glass_thickness = int(
                    current_heliostat_glass_thickness_str
                )
                current_heliostat_area = int(current_heliostat_area_str)
                current_heliostat_l_ratio_str, current_heliostat_w_ratio_str = (
                    current_heliostat_l_w_ratio_str.split("_")
                )
                current_heliostat_l_ratio = int(current_heliostat_l_ratio_str)
                current_heliostat_w_ratio = int(current_heliostat_w_ratio_str)
                current_id_str, _ = os.path.basename(current_path).split(".")
                # [定日镜面积]/[定日镜长宽比]/[定日镜玻璃厚度]/{current_id}.json
                current_id = int(current_id_str)

                with open(current_path, "r") as json_file:
                    json_object = json.load(json_file)
                    # 绘制的是正方形的flux_map
                    current_flux_map_shape = json_object["flux_map_shape"][0]
                    current_distance_heliostat_receiver = json_object[
                        "distance_heliostat_receiver"
                    ]
                    current_sun_elevation = json_object["sun_elevation"]
                    current_sun_azimuth = json_object["sun_azimuth"]
                    current_heliostat_glass_sigma = json_object["heliostat_glass_sigma"]
                    current_sun_csr = json_object["sun_csr"]
                    current_heliostat_glass_refractive_index = json_object[
                        "heliostat_glass_refractive_index"
                    ]
                    current_peak_flux_param = json_object["peak_flux_param"]
                    current_sigma_x = json_object["sigmas"][0]
                    current_sigma_y = json_object["sigmas"][1]
                    current_x13 = json_object["UNIZAR"]["x13"]
                    current_x10 = json_object["UNIZAR"]["x10"]
                    current_y11 = json_object["UNIZAR"]["y11"]
                    current_y10 = json_object["UNIZAR"]["y10"]
                    current_gradient_inv = json_object["UNIZAR"]["gradient_inv"]

                id_array = np.concatenate((id_array, [current_id]), dtype=np.int64)
                heliostat_area_array = np.concatenate(
                    (heliostat_area_array, [current_heliostat_area]), dtype=np.int64
                )
                heliostat_l_ratio_array = np.concatenate(
                    (heliostat_l_ratio_array, [current_heliostat_l_ratio]),
                    dtype=np.int64,
                )
                heliostat_w_ratio_array = np.concatenate(
                    (heliostat_w_ratio_array, [current_heliostat_w_ratio]),
                    dtype=np.int64,
                )
                heliostat_glass_thickness_array = np.concatenate(
                    (
                        heliostat_glass_thickness_array,
                        [current_heliostat_glass_thickness],
                    ),
                    dtype=np.int64,
                )
                flux_map_shape_array = np.concatenate(
                    (flux_map_shape_array, [current_flux_map_shape]), dtype=np.int64
                )
                distance_heliostat_receiver_array = np.concatenate(
                    (
                        distance_heliostat_receiver_array,
                        [current_distance_heliostat_receiver],
                    ),
                    dtype=np.int64,
                )
                sun_elevation_array = np.concatenate(
                    (sun_elevation_array, [current_sun_elevation]), dtype=np.int64
                )
                sun_azimuth_array = np.concatenate(
                    (sun_azimuth_array, [current_sun_azimuth]), dtype=np.int64
                )
                heliostat_glass_sigma_array = np.concatenate(
                    (heliostat_glass_sigma_array, [current_heliostat_glass_sigma]),
                    dtype=np.float64,
                )
                sun_csr_array = np.concatenate(
                    (sun_csr_array, [current_sun_csr]), dtype=np.float64
                )
                heliostat_glass_refractive_index_array = np.concatenate(
                    (
                        heliostat_glass_refractive_index_array,
                        [current_heliostat_glass_refractive_index],
                    ),
                    dtype=np.float64,
                )
                peak_flux_param_array = np.concatenate(
                    (peak_flux_param_array, [current_peak_flux_param]), dtype=np.float64
                )
                sigma_x_array = np.concatenate(
                    (sigma_x_array, [current_sigma_x]), dtype=np.float64
                )
                sigma_y_array = np.concatenate(
                    (sigma_y_array, [current_sigma_y]), dtype=np.float64
                )
                x13_array = np.concatenate((x13_array, [current_x13]), dtype=np.float64)
                x10_array = np.concatenate((x10_array, [current_x10]), dtype=np.float64)
                y11_array = np.concatenate((y11_array, [current_y11]), dtype=np.float64)
                y10_array = np.concatenate((y10_array, [current_y10]), dtype=np.float64)
                gradient_inv_array = np.concatenate(
                    (gradient_inv_array, [current_gradient_inv]), dtype=np.float64
                )

    return (
        id_array,
        heliostat_area_array,
        heliostat_l_ratio_array,
        heliostat_w_ratio_array,
        heliostat_glass_thickness_array,
        flux_map_shape_array,
        distance_heliostat_receiver_array,
        sun_elevation_array,
        sun_azimuth_array,
        heliostat_glass_sigma_array,
        sun_csr_array,
        heliostat_glass_refractive_index_array,
        peak_flux_param_array,
        sigma_x_array,
        sigma_y_array,
        x13_array,
        x10_array,
        y11_array,
        y10_array,
        gradient_inv_array,
    )


def gen_hdf5_file(
    id_array: np.ndarray,
    heliostat_area_array: np.ndarray,
    heliostat_l_ratio_array: np.ndarray,
    heliostat_w_ratio_array: np.ndarray,
    heliostat_glass_thickness_array: np.ndarray,
    flux_map_shape_array: np.ndarray,
    distance_heliostat_receiver_array: np.ndarray,
    sun_elevation_array: np.ndarray,
    sun_azimuth_array: np.ndarray,
    heliostat_glass_sigma_array: np.ndarray,
    sun_csr_array: np.ndarray,
    heliostat_glass_refractive_index_array: np.ndarray,
    peak_flux_param_array: np.ndarray,
    sigma_x_array: np.ndarray,
    sigma_y_array: np.ndarray,
    x13_array: np.ndarray,
    x10_array: np.ndarray,
    y11_array: np.ndarray,
    y10_array: np.ndarray,
    gradient_inv_array: np.ndarray,
):

    h5_file = h5py.File("unizar_fitting_params.hdf5", "w")
    h5_file.create_dataset("id", data=id_array.astype(np.uint32))
    h5_file.create_dataset(
        "heliostat_area", data=heliostat_area_array.astype(np.uint32)
    )
    h5_file.create_dataset(
        "heliostat_l_ratio", data=heliostat_l_ratio_array.astype(np.uint32)
    )
    h5_file.create_dataset(
        "heliostat_w_ratio", data=heliostat_w_ratio_array.astype(np.uint32)
    )
    h5_file.create_dataset(
        "heliostat_glass_thickness",
        data=heliostat_glass_thickness_array.astype(np.uint32),
    )
    h5_file.create_dataset(
        "flux_map_shape", data=flux_map_shape_array.astype(np.uint32)
    )
    h5_file.create_dataset(
        "distance_heliostat_receiver",
        data=distance_heliostat_receiver_array.astype(np.uint32),
    )
    h5_file.create_dataset("sun_elevation", data=sun_elevation_array.astype(np.uint32))
    h5_file.create_dataset("sun_azimuth", data=sun_azimuth_array.astype(np.uint32))
    h5_file.create_dataset(
        "heliostat_glass_sigma", data=heliostat_glass_sigma_array.astype(np.float32)
    )
    h5_file.create_dataset("sun_csr", data=sun_csr_array.astype(np.float32))
    h5_file.create_dataset(
        "heliostat_glass_refractive_index",
        data=heliostat_glass_refractive_index_array.astype(np.float32),
    )
    h5_file.create_dataset(
        "peak_flux_param", data=peak_flux_param_array.astype(np.float32)
    )
    h5_file.create_dataset("sigma_x", data=sigma_x_array.astype(np.float32))
    h5_file.create_dataset("sigma_y", data=sigma_y_array.astype(np.float32))
    h5_file.create_dataset("x13", data=x13_array.astype(np.float32))
    h5_file.create_dataset("x10", data=x10_array.astype(np.float32))
    h5_file.create_dataset("y11", data=y11_array.astype(np.float32))
    h5_file.create_dataset("y10", data=y10_array.astype(np.float32))
    h5_file.create_dataset("gradient_inv", data=gradient_inv_array.astype(np.float32))
    h5_file.close()


def main():
    gen_hdf5_file(*collect_params("./"))


if __name__ == "__main__":
    main()
